import torch                    # 导入 PyTorch 库
from langchain.document_loaders.pdf import PyMuPDFLoader    # 导入 PDF 文档加载器
import re                       # 导入正则表达式模块
from langchain.text_splitter import RecursiveCharacterTextSplitter    # 导入文本分割器
import pdfplumber

#PDF文件的地址
pdf_file_address='./chat_with_rag/Knowledge_db/stuhandbook.pdf'
#Chroma数据库的地址
chroma_db_address="./chat_with_rag/Database/create_db.py" 

#数据读取and数据清洗
def process_pdf(file_path):
    # 创建一个 PyMuPDFLoader Class 实例，输入为待加载的 pdf 文档路径
    loader = PyMuPDFLoader(file_path)

    # 调用 PyMuPDFLoader Class 的函数 load 对 pdf 文件进行加载
    pdf_pages = loader.load()

    #使用正则表达式匹配并删除掉
    pattern = re.compile(r'[^\u4e00-\u9fff](\n)[^\u4e00-\u9fff]', re.DOTALL)
    pdf_pages.page_content = re.sub(pattern, lambda match: match.group(0).replace('', ''), pdf_pages.page_content)
    print(pdf_pages.page_content)
    #去除掉•与空格
    pdf_pages.page_content = pdf_pages.page_content.replace('•', '')
    pdf_pages.page_content = pdf_pages.page_content.replace(' ', '')
    print(pdf_pages.page_content)
    #文档分割
    ''' 
    * RecursiveCharacterTextSplitter 递归字符文本分割
    RecursiveCharacterTextSplitter 将按不同的字符递归地分割(按照这个优先级["", "", " ", ""])，
        这样就能尽量把所有和语义相关的内容尽可能长时间地保留在同一位置
    RecursiveCharacterTextSplitter需要关注的是4个参数：

    * separators - 分隔符字符串数组
    * chunk_size - 每个文档的字符数量限制
    * chunk_overlap - 两份文档重叠区域的长度
    * length_function - 长度计算函数
    '''
    #导入文本分割器
    CHUNK_SIZE = 500

    # 知识库中相邻文本重合长度
    OVERLAP_SIZE = 50
    # 使用递归字符文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=OVERLAP_SIZE
    )
    split_texts = text_splitter.split_text(pdf_pages.page_content[0:1000])
    return split_texts

#表格读取
def extract_tables_from_pdf(pdf_path):
    tables = []
    with pdfplumber.open(pdf_path) as pdf:
        for page in pdf.pages:
            table = page.extract_table()
            if table:
                tables.append(table)
    return tables


#初始化 Chroma 数据库
from chromadb import Chroma                                                     #
from chromadb.config import Settings

# 创建或连接到现有的 Chroma 数据库实例
#连接到一个使用 DuckDB 和 Parquet 格式的 Chroma 数据库
chroma_client = Chroma(settings=Settings(
    chroma_db_impl="duckdb+parquet",
    persist_directory=chroma_db_address                            
))

#加载预训练模型：使用 transformers 库加载一个预训练的模型和分词器
#使用 BERT 或其他适用于文本向量的模型。
from transformers import AutoTokenizer, AutoModel
# 加载预训练的模型和分词器，将文本数据通过分词器处理，然后使用模型生成嵌入向量
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 加载您的文本数据
texts = process_pdf(pdf_file_address)
texts += extract_tables_from_pdf(pdf_file_address)                    

# 文本处理和向量生成
embeddings = []
metadatas = []
ids = []
for text in texts:
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        output = model(**inputs)
    # 提取[CLS]标记的嵌入向量
    embedding = output.last_hidden_state[:, 0, :].numpy()
    embeddings.append(embedding)
    
    # 保存文本的元数据（例如ID或标签）
    metadatas.append({"text": text})
    ids.append(hash(text))  # 使用文本的哈希值作为 ID

# 将文本向量、元数据和 ID 合并

#将向量添加到 Chroma 数据库
# 创建一个"my_collection"的集合
collection_name = "my_collection"
chroma_client.create_collection(name=collection_name)

# 向集合中插入数据，添加文本向量、元数据和 ID 到集合
chroma_client.insert(
    collection_name=collection_name,
    embeddings=embeddings,
    metadatas=metadatas,
    ids=ids
)





